import cv2
from paddleocr import PaddleOCR, draw_ocr
import numpy as np

# 指定模型路径和配置文件（请根据您的实际路径进行修改）
det_model_dir = 'ch_PP-OCRv4_det'  # 检测模型目录
rec_model_dir = 'ch_PP-OCRv4_rec_infer'  # 识别模型目录

# 初始化OCR模型，使用指定的检测和识别模型
ocr = PaddleOCR(
    use_gpu=False,  # 是否使用GPU进行推理，如果没有GPU请设置为False
    det_model_dir=det_model_dir,
    rec_model_dir=rec_model_dir,
    use_angle_cls=True,  # 如果您的模型支持方向分类，请设置为True
    lang='ch'  # 语言类型，对于中文OCR，设置为'ch'
)

# 指定要检测的图像路径（请修改为您的图像路径）
image_path = 'img.png'

# 读取图像
img = cv2.imread(image_path)

# 进行OCR推理
result = ocr.ocr(img, cls=True)

# 绘制检测结果并显示图像（可选）
from PIL import Image

image = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
boxes = [line[0] for line in result]
txts = [line[1][0] for line in result]
scores = [line[1][1] for line in result]
im_show = draw_ocr(image, boxes, txts, scores, font_path='./doc/fonts/simfang.ttf')
im_show = Image.fromarray(cv2.cvtColor(np.asarray(im_show), cv2.COLOR_RGB2BGR))
cv2.imshow('OCR Result', im_show)
cv2.waitKey(0)
cv2.destroyAllWindows()

# 打印检测结果
for line in result:
    line_text = ' '.join([word_info[-1] for word_info in line[1]])
    print(f"Detected text: '{line_text}' with confidence {line[1][0][1]}")